import socket
import datetime
import taosrest
import re


def force_clean_all_databases():
    try:
        # 查询所有数据库
        result = cursor.query("SHOW DATABASES")
        databases = [row[0] for row in result]
        print(f"🧹 找到数据库列表: {databases}")

        for db in databases:
            if db not in ("information_schema", "log", "metrics"):  # 保留系统库
                try:
                    cursor.execute(f"DROP DATABASE {db}")
                    print(f"🗑️ 已删除数据库: {db}")
                except Exception as e:
                    print(f"⚠️ 删除数据库 {db} 失败: {e}")
    except Exception as e:
        print(f"❌ 查询数据库列表失败: {e}")


def create_subtable(source_ip):
    try:
        table_name = f"syslogdb.syslog_logs_{source_ip.replace('.', '_')}"
        cursor.execute(f"""
        CREATE TABLE IF NOT EXISTS {table_name}
        USING syslogdb.syslog_logs TAGS ('generic')
        """)
        print(f"✅ 子表 {table_name} 创建完成或已存在")
    except Exception as e:
        print(f"❌ 子表 {table_name} 创建失败: {e}")


try:
    conn = taosrest.connect(
        url="http://192.168.1.251:6041",
        user="root",
        password="taosdata"
    )
    cursor = conn.cursor()

    # 1. 暴力清空所有用户数据库
    force_clean_all_databases()

    # 2. 创建新的 syslogdb
    cursor.execute("CREATE DATABASE IF NOT EXISTS syslogdb")

    # 3. 创建超级表
    cursor.execute("""
    CREATE STABLE IF NOT EXISTS syslogdb.syslog_logs (
        ts TIMESTAMP,
        source_ip NCHAR(64),
        log_level INT,
        full_message NCHAR(4096)
    ) TAGS (source_type NCHAR(64))
    """)

    print("✅ 连接并初始化数据库成功，开始监听 Syslog...")
except Exception as e:
    print(f"❌ 数据库连接或初始化失败: {e}")
    exit(1)


class SyslogUDPHandler:
    def __init__(self, ip='0.0.0.0', port=514):
        self.server_address = (ip, port)
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(self.server_address)
        print(f"📡 Syslog服务器正在监听 {ip}:{port}")

    def serve_forever(self):
        while True:
            data, addr = self.sock.recvfrom(4096)
            self.handle(addr, data)

    def handle(self, client_address, data):
        try:
            message = data.decode('utf-8', errors='ignore').strip()
            match = re.search(r"<(\d+)>", message)
            priority = int(match.group(1)) if match else 190
            log_level = priority % 8
            source_ip = client_address[0]
            timestamp = datetime.datetime.now().isoformat()

            print(f"📥 收到来自 {source_ip} 的消息: {message}")

            table_name = f"syslogdb.syslog_logs_{source_ip.replace('.', '_')}"

            # 子表确保存在
            create_subtable(source_ip)

            # 转义单引号
            escaped_message = message.replace("'", "''")

            # 插入数据
            cursor.execute(f"""
            INSERT INTO {table_name} (ts, source_ip, log_level, full_message)
            VALUES ('{timestamp}', '{source_ip}', {log_level}, '{escaped_message}')
            """)

        except Exception as e:
            print(f"⚠️ 插入失败: {e}")


if __name__ == "__main__":
    server = SyslogUDPHandler()
    server.serve_forever()